
#include "mysql_protocol.h"
#include "tds_protocol.h"
#include "log.h"
namespace secsmart_protocol
{
    DETECT_RESULT MySQLProtocol::parse(secsmart_session::session_sptr session, const uint8_t *buf, size_t size, bool is_forward)
    {
        direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_UNKNOW;
        if (size < 4){ //小于协议长度
            return DETECT_RESULT::DETECT_UNKNOWN;
        }

        size_t pos = 0;
        if (left_length_ > 0)
        {
            pos += left_length_;
            left_length_ = 0;
        }
        
        if(pos + 2 >= size){
            return DETECT_RESULT::DETECT_UNKNOWN;
        }

        uint32_t len = ((uint8_t)(buf[pos + 2]) << 16) + ((uint8_t)(buf[pos + 1]) << 8) + (uint8_t)(buf[pos]) + 4;
        if(len > size){
            addNoMatchCount();
            return DETECT_RESULT::DETECT_UNKNOWN;
        }

        if (len == size)
        {
            uint8_t packet_number = (uint8_t)buf[3];
            if(size > 6)
            {
                uint8_t server_version = (uint8_t)buf[5];
                if (packet_number == 0 && server_version > '0' && server_version < '9' && buf[6] == '.')
                {
                    direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_RESPONSE;
                    if(!setIpPort(session, getDirction(), is_forward)){
                        LOG_DEBUG("set port ip fail");
                        return DETECT_RESULT::DETECT_FAIL;
                    }
                    phase_ = PHASE_FULL_LINK;
                    parseVersion(buf + 5, size - 5);
                    addMatchCount();
                    return DETECT_RESULT::DETECT_OK;
                }
            }
            
            if (phase_ == PHASE_FULL_LINK && packet_number == 1 && !gotLoginReq_)
            {
                direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_REQUEST;
                LOG_DEBUG("MySQL : Login Request Packet detected!");
                gotLoginReq_ = true;
                addMatchCount();
                return DETECT_RESULT::DETECT_OK;
            }

            uint8_t flag = (uint8_t)buf[4];

            // Response Packet
            if (flag == 0 || flag == 0xFE || flag == 0xFF)
            {
                LOG_DEBUG("MySQL : Response Packet detected!");
                direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_RESPONSE;
                if(!setIpPort(session, getDirction(), is_forward)){
                    LOG_DEBUG("set port ip fail");
                    return DETECT_RESULT::DETECT_FAIL;
                }
                if(flag == 0){ //识别binlog
                  if((uint8_t)buf[5] == 0xef){
                    ++binlogSize;
                  }else{
                    isBinlogResponse(buf, size);
                  }
                }
                addMatchCount();
                return DETECT_RESULT::DETECT_OK;
            }

            // Command Packet
            if (flag == 1 && len ==5)
            {
                // COM_QUIT
                LOG_DEBUG("MySQL : Command Packet detected!");
                direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_REQUEST;
                if(!setIpPort(session, getDirction(), is_forward)){
                    LOG_DEBUG("set port ip fail");
                    return DETECT_RESULT::DETECT_FAIL;
                }
                addMatchCount();
                return DETECT_RESULT::DETECT_OK;
            }

            // Query Packet
            if (flag == 3)
            {
                // 检查是否为gbase
                if (buf[8] == 0x67 && buf[9] == 0x62 && buf[10] == 0x61 && buf[11] == 0x73 && buf[12] == 0x65)
                {
                    LOG_DEBUG("GBASE8A : Query Packet detected!");
                    databaseName = DBNAME_GBASE8A;
                    direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_REQUEST;
                    if(!setIpPort(session, getDirction(), is_forward)){
                        LOG_DEBUG("set port ip fail");
                        return DETECT_RESULT::DETECT_FAIL;
                    }
                    addMatchCount();
                    return DETECT_RESULT::DETECT_OK;
                }
                LOG_DEBUG("MySQL : Query Packet detected!");
            }

            if ((flag > 1 && flag < 0x1f) || flag == 0xef) //0xef binlog
            {
                direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_REQUEST;
                if(!setIpPort(session, getDirction(), is_forward)){
                    LOG_DEBUG("set port ip fail");
                    return DETECT_RESULT::DETECT_FAIL;
                }
                if((flag == 0xef || flag == 0x12) && packet_number == 0){
                    ++binlogSize;
                }
                addMatchCount();
                return DETECT_RESULT::DETECT_OK;
            }
            addNoMatchCount();
            return DETECT_RESULT::DETECT_UNKNOWN;
        }

        // len < size 的情况

        uint32_t last_packet_number = buf[pos + 3];

        // Check TABULAR Resp
        pos += len;

        size_t total_len = pos;
        size_t nLen = size - total_len;

        uint32_t packet_number = packet_number_ + 1;
        while (nLen > 0 && pos + 4 < size)
        {
            size_t pdu_size = ((uint8_t)(buf[pos + 2]) << 16) + ((uint8_t)(buf[pos + 1]) << 8) + (uint8_t)(buf[pos]) + 4;
            last_packet_number = buf[pos + 3];
            if (pdu_size > size)
            {
                break;
            }
            
            nLen -= pdu_size;
            pos += pdu_size;
            total_len += pdu_size;

            packet_number = (packet_number == 255) ? 0 : packet_number + 1;
        }

        if (total_len == size)
        {
            left_length_ = 0;
            packet_number_ = last_packet_number;
            LOG_DEBUG("MySQL : TABULAR Resp detected!");
            direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_RESPONSE;
            if(!setIpPort(session, getDirction(), is_forward)){
                LOG_DEBUG("set port ip fail");
                return DETECT_RESULT::DETECT_FAIL;
            }
            if(size > 6 && buf[4] == 0 && (uint8_t)buf[5] == 0xef)
            {
                ++binlogSize;
            }else{
                isBinlogResponse(buf, size);
            }
            addMatchCount();
            return DETECT_RESULT::DETECT_OK;
        }

        // Packet data segment,Check Packet Number
        if (total_len > size)
        {
            left_length_ = total_len - size;
            if (packet_number == last_packet_number)
            {
                LOG_DEBUG("MySQL : Packet data segment detected!");
                packet_number_ = last_packet_number;
                direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_RESPONSE;
                if(!setIpPort(session, getDirction(), is_forward)){
                    LOG_DEBUG("set port ip fail");
                    return DETECT_RESULT::DETECT_FAIL;
                }
                addMatchCount();
                return DETECT_RESULT::DETECT_OK;
            }
        }

        left_length_ = 0;
        packet_number_ = 0;

        if (len + 3 == size)
        {
            if (checkIfTheDataIsCompressed(buf, size))
            {
                LOG_DEBUG("MySQL : Compressed Packet detected!");
                if (last_packet_number == 0)
                {
                    direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_REQUEST;
                    if(!setIpPort(session, getDirction(), is_forward)){
                        LOG_DEBUG("set port ip fail");
                        return DETECT_RESULT::DETECT_FAIL;
                    }
                    addMatchCount();
                    return DETECT_RESULT::DETECT_OK;
                }
                else if (last_packet_number == 1)
                {
                    direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_RESPONSE;
                    if(!setIpPort(session, getDirction(), is_forward)){
                        LOG_DEBUG("set port ip fail");
                        return DETECT_RESULT::DETECT_FAIL;
                    }
                    addMatchCount();
                    return DETECT_RESULT::DETECT_OK;
                }
            }
        }
    
        addNoMatchCount();
        return DETECT_RESULT::DETECT_UNKNOWN;

    }

    bool MySQLProtocol::checkIfTheDataIsCompressed(const uint8_t *buf, size_t size)
    {
        if(size < 10){
            return false;
        }
        unsigned long tmp_pktLen = 0;
        unsigned long tmp_compBufLen = (buf[2] << 16) | (buf[1] << 8) | buf[0];
        if (0 == buf[4] && 0 == buf[5] && 0 == buf[6])
        {
            // 压缩协议开启，但数据未进行压缩
            tmp_pktLen = (buf[9] << 16) | (buf[8] << 8) | (buf[7]);
        }
        return (tmp_compBufLen == (tmp_pktLen + 4));
    }


    void MySQLProtocol::parseVersion(const uint8_t *buf, size_t size)
    {
        versionInfo.clear();
        for(size_t i = 0; i < size && buf[i] != '\0'; ++i){
            versionInfo += buf[i];
        }
        version = versionInfo;
        if(version.find("MariaDB") != std::string::npos){
            databaseName = DBNAME_MARIADB;
        }else  if(version.find("TiDB") != std::string::npos){
            databaseName = DBNAME_TIDB;
        }

    }

    void MySQLProtocol::isBinlogResponse(const uint8_t *buf, size_t size)
    {
        LOG_DEBUG("MySQL : size = {}", size);
        if(size < 25){
            return;
        }
        uint32_t eventSize = ((uint8_t)(buf[17]) << 24) + ((uint8_t)(buf[16]) << 16) + ((uint8_t)(buf[15]) << 8) + (uint8_t)(buf[14]);
        uint32_t packetSize =((uint8_t)(buf[2]) << 16) + ((uint8_t)(buf[1]) << 8) + (uint8_t)(buf[0]);
        LOG_DEBUG("MySQL : eventSize = {} ,packetSize = {}", eventSize, packetSize);
        if(eventSize + 1 != packetSize)
        {
            return;
        }
        uint16_t eventFlags= ((uint8_t)(buf[23]) << 8) + (uint8_t)(buf[22]);
        LOG_DEBUG("MySQL : eventFlags = {}", (int)eventFlags);
        switch(eventFlags)
        {
            case 0x0000:
            case 0x0004:
            case 0x0008:
            case 0x0020:
            case 0x0040:
            case 0x0080:
            case 0x0100:
            case 0x0200:  
                break;
            default:
                return;

        }
        
        uint8_t eventType= (uint8_t)(buf[9]);
        LOG_DEBUG("MySQL : eventType = {}", (int)eventType);
        if((eventType >= 0 && eventType < 0x24) || eventType == 0x27)
        {
            ++binlogSize;
        }
        return;
    }

    std::string MySQLProtocol::extendInfo(){
        LOG_DEBUG("MySQL : binlogSize = {}", binlogSize);
        if(binlogSize > 1){
            return ", \"binlog\" : true";
        }
        return "";
    }

}